package utils

import java.io.File

import config.MyFile
import org.apache.commons.io.FileUtils

import scala.collection.mutable
import scala.jdk.CollectionConverters._

object GetGene extends MyFile {


  def main(args: Array[String]): Unit = {

    val path = "D:\\黄瓜和南瓜信息数据库开发\\黄瓜和南瓜数据\\2021.5.24更新要求\\Cm"
      getGene(s"$path/Cmoschata_v1.gff3",s"$path/Cmoschata_v1.chr.fa",s"$path/Cmoschata_gene.fa")

    val path2 = "D:\\黄瓜和南瓜信息数据库开发\\黄瓜和南瓜数据\\2021.5.24更新要求\\Cs"
    getGene(s"$path2/Gy14_v2.gene.gff", s"$path2/Gy14_v2.0_final.fa", s"$path2/Gy14_v2.0_gene.fa")

  }

  /**
   * 得到CDS和 PEP可以用gffread软件（cufflinks）
   */

  def getGene(gffPath: String, fastaPath: String, outPath: String) = {


    val gff = gffPath.readLines
    val fasta = fastaPath.readFileToString

    val fa = fasta.split(">").tail.map{x=>
      val line = x.split("\n")
      line.head -> line.tail.mkString
    }.toMap

    val atcg = Map("A" -> "T", "C" -> "G", "N" -> "N", "T" -> "A", "G" -> "C",
      "R" -> "Y", "Y" -> "R", "M" -> "K", "K" -> "M",
      "S" -> "S", "W" -> "W", "H" -> "D", "D" -> "H", "B" -> "V", "V" -> "B")

    val genes = gff.filter(!_.startsWith("#")).map(_.split("\t")).
      filter(x => x(2) == "gene").map { x =>
      val id = x.last.split(";").head.drop(3)
      val start = x(3).toInt
      val end = x(4).toInt
      val strand = x(6)
      val chr = x.head
      val fas = fa(chr)

      val gene = fas.slice(start - 1, end)
      val seq = if (strand == "+") {
        gene
      } else {
        gene.reverse.map { y =>
            atcg(y.toUpper.toString)
        }.mkString
      }
      (s">$id", seq)
    }.groupBy(_._1).flatMap { y =>
      if (y._2.length == 1) {
        y._2.flatMap(z => mutable.Buffer(z._1.replaceAll(" ", "_"), z._2))
      } else {
        y._2.zipWithIndex.flatMap { z =>
          val geneids = z._1._1
          val geneid = geneids.replaceAll(" ", "_") + "_" + (z._2 + 1)
          mutable.Buffer(s"$geneid", z._1._2)
        }
      }

    }.toBuffer
    FileUtils.writeLines(new File(outPath), genes.asJava)
  }


}
